#!/usr/bin/env python3
import glob
import json
import os
import shutil

import numpy as np


def make_calypso_input(
    nameofatoms,
    numberofatoms,
    numberofformula,
    volume,
    distanceofion,
    psoratio,
    popsize,
    maxstep,
    icode,
    split,
    vsc,
    maxnumatom,
    ctrlrange,
    pstress,
    fmax,
):
    ret = "################################ The Basic Parameters of CALYPSO ################################\n"
    ret += "# A string of one or several words contain a descriptive name of the system (max. 40 characters).\n"
    assert nameofatoms is not None
    ret += "SystemName = %s\n" % ("".join(nameofatoms))
    ret += "# Number of different atomic species in the simulation.\n"
    ret += "NumberOfSpecies = %d\n" % (len(nameofatoms))
    ret += "# Element symbols of the different chemical species.\n"
    ret += "NameOfAtoms = %s\n" % (" ".join(nameofatoms))
    ret += "# Number of atoms for each chemical species in one formula unit. \n"
    assert numberofatoms is not None and len(numberofatoms) == len(nameofatoms)
    ret += "NumberOfAtoms = %s\n" % (" ".join(list(map(str, numberofatoms))))
    ret += "# The range of formula unit per cell in your simulation. \n"
    assert (
        numberofformula is not None
        and len(numberofformula) == 2
        and isinstance(numberofformula, list)
    )
    ret += "NumberOfFormula = %s\n" % (" ".join(list(map(str, numberofformula))))
    ret += "# The volume per formula unit. Unit is in angstrom^3.\n"
    if volume is None:
        ret += "# volume not found, CALYPSO will set one!\n"
    else:
        ret += "Volume = %s\n" % (volume)
    ret += "# Minimal distance between atoms of each chemical species. Unit is in angstrom.\n"
    assert len(distanceofion) == len(
        nameofatoms
    )  # "check distance of ions and the number of atoms"
    assert len(distanceofion[0]) == len(nameofatoms)
    ret += "@DistanceOfIon \n"
    for temp in distanceofion:
        ret += "%4s \n" % (" ".join(list(map(str, temp))))
    ret += "@End\n"
    ret += "# It determines which algorithm should be adopted in the simulation.\n"
    ret += "Ialgo = 2\n"
    ret += "# Ialgo = 1 for Global PSO\n"
    ret += "# Ialgo = 2 for Local PSO (default value)\n"
    ret += "# The proportion of the structures generated by PSO.\n"
    assert 0 <= psoratio <= 1
    ret += "PsoRatio = %s\n" % (psoratio)
    ret += (
        "# The population size. Normally, it has a larger number for larger systems.\n"
    )
    assert popsize is not None and isinstance(popsize, int)
    ret += "PopSize = %d\n" % (popsize)
    assert maxstep is not None and isinstance(maxstep, int)
    ret += "# The Max step for iteration\n"
    ret += "MaxStep = %d\n" % (maxstep)
    ret += "#It determines which method should be adopted in generation the random structure. \n"
    ret += "GenType= 1 \n"
    ret += "# 1 under symmetric constraints\n"
    ret += "# 2 grid method for large system\n"
    ret += "# 3 and 4 core grow method \n"
    ret += "# 0 combination of all method\n"
    ret += "# If GenType=3 or 4, it determined the small unit to grow the whole structure\n"
    ret += "# It determines which local optimization method should be interfaced in the simulation.\n"
    assert icode is not None and isinstance(icode, int)
    ret += "ICode= %d\n" % (icode)
    ret += "# ICode= 1 interfaced with VASP\n"
    ret += "# ICode= 2 interfaced with SIESTA\n"
    ret += "# ICode= 3 interfaced with GULP\n"
    ret += "# The number of lbest for local PSO\n"
    ret += "NumberOfLbest=4\n"
    ret += "# The Number of local optimization for each structure.\n"
    ret += "NumberOfLocalOptim= 3\n"
    ret += "# The command to perform local optimiztion calculation (e.g., VASP, SIESTA) on your computer.\n"
    ret += "Command = sh submit.sh\n"
    ret += "MaxTime = 9000 \n"
    ret += "# If True, a previous calculation will be continued.\n"
    ret += "PickUp = F\n"
    ret += "# At which step will the previous calculation be picked up.\n"
    ret += "PickStep = 1\n"
    ret += "# If True, the local optimizations performed by parallel\n"
    ret += "Parallel = F\n"
    ret += "# The number node for parallel \n"
    ret += "NumberOfParallel = 4\n"
    assert split is not None
    ret += "Split = %s\n" % (split)
    assert pstress is not None and (
        isinstance(pstress, int) or isinstance(pstress, float)
    )
    ret += "PSTRESS = %f\n" % (pstress)
    assert fmax is not None or isinstance(fmax, float)
    ret += "fmax = %f\n" % (fmax)
    ret += "################################ End of The Basic Parameters of CALYPSO #######################\n"
    if vsc == "T":
        assert len(ctrlrange) == len(
            nameofatoms
        )  #'check distance of ions and the number of atoms'
        ret += "##### The Parameters For Variational Stoichiometry  ##############\n"
        ret += (
            "## If True, Variational Stoichiometry structure prediction is performed\n"
        )
        ret += "VSC = %s\n" % (vsc)
        ret += "# The Max Number of Atoms in unit cell\n"
        ret += "MaxNumAtom = %s\n" % (maxnumatom)
        ret += "# The Variation Range for each type atom \n"
        ret += "@CtrlRange\n"
        for ttemp in ctrlrange:
            ret += "%4s \n" % (" ".join(list(map(str, ttemp))))
        ret += "@end\n"
        ret += "###################End Parameters for VSC ##########################\n"
    return ret


def _make_model_devi_buffet(jdata, calypso_run_opt_path):
    calypso_input_path = jdata.get("calypso_input_path")
    if jdata.get("vsc", False):
        # [input.dat.Li.250, input.dat.Li.300]
        one_ele_inputdat_list = list(
            set(
                glob.glob(
                    f"{jdata.get('calypso_input_path')}/input.dat.{jdata.get('type_map')[0]}.*"
                )
            )
        )
        # [input.dat.La, input.dat.H, input.dat.LaH,] only one pressure
        if len(one_ele_inputdat_list) == 0:
            os.system(f"cp {calypso_input_path}/input.dat.* {calypso_run_opt_path[0]}")
        # different pressure, 250GPa and 300GPa
        # [input.dat.La.250, input.dat.H.250, input.dat.LaH.250, input.dat.La.300, input.dat.H.300, input.dat.LaH.300,]
        else:
            pressures_list = [temp.split(".")[-1] for temp in one_ele_inputdat_list]
            pressures_list = list(map(int, pressures_list))
            # calypso_run_opt_path = ['gen_struc_analy.000','gen_struc_analy.001']
            for press_idx, temp_calypso_run_opt_path in enumerate(calypso_run_opt_path):
                cur_press = pressures_list[press_idx]
                os.system(
                    f"cp {calypso_input_path}/input.dat.*.{cur_press} {temp_calypso_run_opt_path}"
                )
    elif not jdata.get("vsc", False):
        shutil.copyfile(
            os.path.join(calypso_input_path, "input.dat"),
            os.path.join(calypso_run_opt_path[0], "input.dat"),
        )
        if not os.path.exists(os.path.join(calypso_run_opt_path[0], "input.dat")):
            raise FileNotFoundError("input.dat")


def _make_model_devi_native_calypso(iter_index, model_devi_jobs, calypso_run_opt_path):
    for iiidx, jobbs in enumerate(model_devi_jobs):
        if iter_index in jobbs.get("times"):
            cur_job = model_devi_jobs[iiidx]

    work_path = os.path.dirname(calypso_run_opt_path[0])
    # cur_job.json
    with open(os.path.join(work_path, "cur_job.json"), "w") as outfile:
        json.dump(cur_job, outfile, indent=4)

    # Crystal Parameters
    nameofatoms = cur_job.get("NameOfAtoms")
    numberofatoms = cur_job.get("NumberOfAtoms")
    numberofformula = cur_job.get("NumberOfFormula", [1, 1])
    volume = cur_job.get("Volume")
    distanceofion = cur_job.get("DistanceOfIon")
    psoratio = cur_job.get("PsoRatio", 0.6)
    popsize = cur_job.get("PopSize", 30)
    maxstep = cur_job.get("MaxStep", 5)
    icode = cur_job.get("ICode", 1)
    split = cur_job.get("Split", "T")
    # Cluster

    # 2D

    # VSC Control
    maxnumatom = None
    ctrlrange = None
    vsc = cur_job.get("VSC", "F")
    if vsc == "T":
        maxnumatom = cur_job.get("MaxNumAtom")
        ctrlrange = cur_job.get("CtrlRange")
    # Optimization
    fmax = cur_job.get("fmax", 0.01)
    # pstress is a List which contains the target stress
    pstress = cur_job.get("PSTRESS", [0.001])
    # pressures
    for press_idx, temp_calypso_run_opt_path in enumerate(calypso_run_opt_path):
        # cur_press
        cur_press = pstress[press_idx]
        file_c = make_calypso_input(
            nameofatoms,
            numberofatoms,
            numberofformula,
            volume,
            distanceofion,
            psoratio,
            popsize,
            maxstep,
            icode,
            split,
            vsc,
            maxnumatom,
            ctrlrange,
            cur_press,
            fmax,
        )
        with open(os.path.join(temp_calypso_run_opt_path, "input.dat"), "w") as cin:
            cin.write(file_c)


def write_model_devi_out(devi, fname):
    assert devi.shape[1] == 8
    # assert devi.shape[1] == 7
    header = "%5s" % "step"
    for item in "vf":
        header += "%16s%16s%16s" % (
            f"max_devi_{item}",
            f"min_devi_{item}",
            f"avg_devi_{item}",
        )
    header += "%16s" % "min_dis"
    np.savetxt(
        fname,
        devi,
        fmt=["%5d"] + ["%17.6e" for _ in range(7)],
        delimiter="",
        header=header,
    )
    return devi
